import os
import tempfile
from io import BytesIO

from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Query
from fastapi.responses import FileResponse, StreamingResponse
from pydub import AudioSegment
from indextts.infer import IndexTTS  # 确保路径正确

app = FastAPI(
    title="IndexTTS语音合成服务",
    description="基于IndexTTS二次开发的语音合成API服务",
    version="1.0.0",
    openapi_tags=[
        {
            "name": "音频生成",
            "description": "核心语音合成接口",
            "externalDocs": {
                "description": "林夕",
                "url": "mailto:2385564035@qq.com"
            }
        }
    ],
)
tts_engine = IndexTTS(model_dir="checkpoints", cfg_path="checkpoints/config.yaml")

SUPPORTED_FORMATS = {
    "wav": "wav",
    "mp3": "mp3",
    "ogg": "ogg",
    "flac": "flac"
}

def cleanup_files(*files):
    """清理临时文件"""
    for file_path in files:
        try:
            if file_path and os.path.exists(file_path):
                os.remove(file_path)
        except Exception as e:
            print(f"Error deleting file {file_path}: {e}")

@app.post(
    "/generate_audio_file",
    tags=["音频生成"],
    summary="生成语音文件",
    description="""根据参考音频克隆声音并生成指定格式的语音文件  
    **典型场景**:
    - 语音克隆合成
    - 多格式音频输出
    """,
    response_description="生成的音频文件流",
    responses={
        200: {
            "content": {"audio/*": {}},
            "description": "成功返回音频文件"
        },
        400: {
            "content": {
                "application/json": {
                    "example": {"detail": "Unsupported format. Available: ['wav', 'mp3']"}
                }
            }
        },
        500: {
            "content": {
                "application/json": {
                    "example": {"detail": "Generation failed: [error message]"}
                }
            }
        }
    }
)
async def generate_audio_file(
    text: str = Query(..., description="需要合成的文本内容，支持中文/英文混合", example="欢迎使用IndexTTS"),
    reference_audio: UploadFile = File(..., description="参考音频文件(WAV格式，时长5-30秒)"),
    format: str = Query("wav", description="输出音频格式", enum=list(SUPPORTED_FORMATS.keys())),
    background_tasks: BackgroundTasks = BackgroundTasks()
):
    """
    生成并返回音频文件
    Parameters:
    - text: 要合成的文本
    - reference_audio: 参考音频文件
    - format: 输出格式 (wav, mp3, ogg, flac)
    """
    if format.lower() not in SUPPORTED_FORMATS:
        raise HTTPException(400, f"Unsupported format. Available: {list(SUPPORTED_FORMATS.keys())}")

    ref_path = None
    output_path = None
    converted_path = None

    try:
        # 保存参考音频到临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_ref:
            contents = await reference_audio.read()
            tmp_ref.write(contents)
            ref_path = tmp_ref.name

        # 生成音频到临时文件
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_out:
            output_path = tmp_out.name

        tts_engine.infer(ref_path, text, output_path)

        final_output = output_path

        # 格式转换
        if format.lower() != "wav":
            audio = AudioSegment.from_wav(output_path)
            base = os.path.splitext(output_path)[0]  # 去除原扩展名
            converted_path = f"{base}.{format}"
            audio.export(converted_path, format=format)
            os.remove(output_path)  # 删除原始wav文件
            final_output = converted_path

        # 添加后台清理任务
        background_tasks.add_task(cleanup_files, ref_path, final_output)

        return FileResponse(
            final_output,
            media_type=f"audio/{format}",
            filename=f"output.{format}"
        )

    except Exception as e:
        # 发生错误时立即清理
        cleanup_files(ref_path, output_path, converted_path)
        raise HTTPException(500, f"Generation failed: {str(e)}")


@app.post(
    "/generate_audio_data",
    tags=["音频生成"],
    summary="生成音频数据流",
    description="""直接返回音频二进制数据流，适用于实时场景  
    **典型场景**:
    - 低延迟语音合成
    - 网页音频即时播放
    - 嵌入式设备语音输出
    """,
    response_description="原始音频数据流",
    responses={
        200: {
            "content": {"audio/*": {}},
            "description": "成功返回音频流",
            "headers": {
                "Content-Disposition": {
                    "description": "下载文件名",
                    "schema": {"type": "string"}
                }
            }
        },
        400: {
            "content": {
                "application/json": {
                    "example": {"detail": "Unsupported format. Available: ['wav', 'mp3']"}
                }
            }
        },
        500: {
            "content": {
                "application/json": {
                    "example": {"detail": "Generation failed: [error message]"}
                }
            }
        }
    }
)
async def generate_audio_data(
        text: str = Query(...,
                          description="需要合成的文本内容，支持中文/英文混合",
                          example="欢迎使用IndexTTS",),
        reference_audio: UploadFile = File(...,
                                           description="参考音频文件(WAV格式，时长5-30秒)",
                                           example="example.wav"),
        format: str = Query("wav",
                            description="输出音频格式",
                            enum=list(SUPPORTED_FORMATS.keys()))
):
    """
    生成并返回音频二进制流

    Parameters:
    - text: 要合成的文本（2-500字符）
    - reference_audio: 参考音频文件(WAV格式)
    - format: 输出格式 (wav, mp3, ogg, flac)
    """
    if format.lower() not in SUPPORTED_FORMATS:
        raise HTTPException(400, f"Unsupported format. Available: {list(SUPPORTED_FORMATS.keys())}")

    ref_path = None
    try:
        # 保存参考音频到临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_ref:
            contents = await reference_audio.read()
            tmp_ref.write(contents)
            ref_path = tmp_ref.name

        # 生成音频到临时文件
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_out:
            output_path = tmp_out.name
            tts_engine.infer(ref_path, text, output_path)

            # 读取并转换格式
            audio = AudioSegment.from_wav(output_path)
            buffer = BytesIO()
            audio.export(buffer, format=format)
            buffer.seek(0)

        return StreamingResponse(
            buffer,
            media_type=f"audio/{format}",
            headers={"Content-Disposition": f"attachment; filename=output.{format}"}
        )

    except Exception as e:
        raise HTTPException(500, f"Generation failed: {str(e)}")
    finally:
        cleanup_files(ref_path, output_path if 'output_path' in locals() else None)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)